"""
Copyright (c) 2022 Beijing Jiaotong University
PhotLab is licensed under [Open Source License].
You can use this software according to the terms and conditions of the [Open Source License].
You may obtain a copy of [Open Source License] at: [https://open.source.license/]

THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.

See the [Open Source License] for more details.

Author: Chunyu Li
Created: 2022/9/6
Supported by: National Key Research and Development Program of China
"""


from ...utils import plot_scatter

import matplotlib.pyplot as plt
import numpy as np


def constellation_diagram(signals, isdata: bool = True, os_name=None):
    """
    绘制星座图
    :param signals : 信号数据 
    :param isdata : 是否需要返回绘图所需数据
    :returns: 
    """
    if len(signals) == 4:
        # prev_symbols = [signals[2], signals[3]]
        signals = [signals[0], signals[1]]
    else:
        pass
    if isdata == False:
        if len(signals) > 1:  # 如果接收的信号的偏振数大于1 则画出所有偏振方向的星座图
            for i in range(len(signals)):
                plot_scatter(
                    signals[i], pt_size=1, title="Constellation diagram", xlabel="In-phase [a.u.]", ylabel="Quadrature [a.u.]", os_name = os_name
                )
        else:
            plot_scatter(
                signals[0], pt_size=1, title="Constellation diagram", xlabel="In-phase [a.u.]", ylabel="Quadrature [a.u.]", 
            )
        # signals.append(prev_symbols[0])
        # signals.append(prev_symbols[1])
    else:
        plot_num = len(signals)
        if len(signals) > 1:   # 如果接收的信号的偏振数大于1 则返回所有方向的绘图所需数据
            axis_x = []
            axis_y = []
            for i in range(len(signals)):
                axis_x.append(np.real(signals[i]).ravel())
                axis_y.append(np.imag(signals[i]).ravel())
            return [axis_x, axis_y], plot_num
        else:
            axis_x = np.real(signals[0]).ravel()
            axis_y = np.imag(signals[0]).ravel()
            return [axis_x, axis_y], plot_num
        # signals.append(prev_symbols[0])
        # signals.append(prev_symbols[1])
